"""
视频处理工具
包含视频读取、写入、格式转换、批量处理等功能
"""

import cv2
import numpy as np
from PIL import Image
import moviepy.editor as mp
from moviepy.video.fx import resize, crop, speedx
import os
from pathlib import Path
import logging
from typing import List, Tuple, Optional, Dict, Any
import json
import time

logger = logging.getLogger(__name__)


class VideoProcessor:
    """视频处理器"""
    
    def __init__(self):
        self.supported_formats = ['.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv', '.webm']
        self.supported_image_formats = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
    
    def get_video_info(self, video_path: str) -> Dict[str, Any]:
        """获取视频信息"""
        try:
            cap = cv2.VideoCapture(video_path)
            
            if not cap.isOpened():
                raise ValueError(f"无法打开视频文件: {video_path}")
            
            info = {
                'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                'fps': cap.get(cv2.CAP_PROP_FPS),
                'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
                'duration': cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS),
                'codec': int(cap.get(cv2.CAP_PROP_FOURCC)),
                'format': Path(video_path).suffix.lower()
            }
            
            cap.release()
            return info
            
        except Exception as e:
            logger.error(f"获取视频信息失败: {e}")
            raise
    
    def read_video_frames(self, video_path: str, start_frame: int = 0, 
                         end_frame: Optional[int] = None) -> List[np.ndarray]:
        """读取视频帧"""
        try:
            cap = cv2.VideoCapture(video_path)
            
            if not cap.isOpened():
                raise ValueError(f"无法打开视频文件: {video_path}")
            
            frames = []
            frame_idx = 0
            
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                
                if frame_idx >= start_frame:
                    # 转换BGR到RGB
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    frames.append(frame_rgb)
                
                frame_idx += 1
                
                if end_frame and frame_idx >= end_frame:
                    break
            
            cap.release()
            return frames
            
        except Exception as e:
            logger.error(f"读取视频帧失败: {e}")
            raise
    
    def save_video_frames(self, frames: List[np.ndarray], output_path: str, 
                         fps: int = 24, codec: str = 'mp4v') -> str:
        """保存帧序列为视频"""
        try:
            if not frames:
                raise ValueError("没有帧可以保存")
            
            # 获取视频尺寸
            height, width = frames[0].shape[:2]
            
            # 创建视频写入器
            fourcc = cv2.VideoWriter_fourcc(*codec)
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            
            # 写入帧
            for frame in frames:
                # 转换RGB到BGR
                frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                out.write(frame_bgr)
            
            out.release()
            logger.info(f"视频保存成功: {output_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"保存视频失败: {e}")
            raise
    
    def extract_frames(self, video_path: str, output_dir: str, 
                      frame_rate: Optional[int] = None, 
                      start_time: float = 0, 
                      end_time: Optional[float] = None) -> List[str]:
        """提取视频帧"""
        try:
            # 创建输出目录
            Path(output_dir).mkdir(parents=True, exist_ok=True)
            
            # 使用moviepy加载视频
            video = mp.VideoFileClip(video_path)
            
            # 裁剪时间范围
            if end_time:
                video = video.subclip(start_time, end_time)
            else:
                video = video.subclip(start_time)
            
            # 提取帧
            frame_paths = []
            frame_count = 0
            
            if frame_rate:
                # 按指定帧率提取
                for t in np.arange(0, video.duration, 1/frame_rate):
                    frame = video.get_frame(t)
                    frame_path = os.path.join(output_dir, f"frame_{frame_count:06d}.jpg")
                    Image.fromarray(frame).save(frame_path)
                    frame_paths.append(frame_path)
                    frame_count += 1
            else:
                # 提取所有帧
                for frame in video.iter_frames():
                    frame_path = os.path.join(output_dir, f"frame_{frame_count:06d}.jpg")
                    Image.fromarray(frame).save(frame_path)
                    frame_paths.append(frame_path)
                    frame_count += 1
            
            video.close()
            logger.info(f"提取了 {len(frame_paths)} 帧到 {output_dir}")
            return frame_paths
            
        except Exception as e:
            logger.error(f"提取帧失败: {e}")
            raise
    
    def create_video_from_frames(self, frame_dir: str, output_path: str, 
                                fps: int = 24, pattern: str = "frame_*.jpg") -> str:
        """从帧序列创建视频"""
        try:
            # 获取所有帧文件
            frame_files = sorted(Path(frame_dir).glob(pattern))
            
            if not frame_files:
                raise ValueError(f"在 {frame_dir} 中没有找到匹配的帧文件")
            
            # 读取第一帧获取尺寸
            first_frame = Image.open(frame_files[0])
            width, height = first_frame.size
            
            # 创建视频写入器
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            
            # 写入帧
            for frame_file in frame_files:
                frame = cv2.imread(str(frame_file))
                out.write(frame)
            
            out.release()
            logger.info(f"从帧序列创建视频成功: {output_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"从帧序列创建视频失败: {e}")
            raise
    
    def convert_format(self, input_path: str, output_path: str, 
                      codec: str = 'mp4v', quality: int = 23) -> str:
        """转换视频格式"""
        try:
            # 使用moviepy转换
            video = mp.VideoFileClip(input_path)
            
            # 确保输出目录存在
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            
            # 写入新格式
            video.write_videofile(
                output_path,
                codec='libx264' if codec == 'mp4v' else codec,
                audio_codec='aac',
                ffmpeg_params=['-crf', str(quality)]
            )
            
            video.close()
            logger.info(f"格式转换成功: {input_path} -> {output_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"格式转换失败: {e}")
            raise
    
    def resize_video(self, input_path: str, output_path: str, 
                    width: Optional[int] = None, 
                    height: Optional[int] = None, 
                    scale: Optional[float] = None) -> str:
        """调整视频尺寸"""
        try:
            video = mp.VideoFileClip(input_path)
            
            # 计算新尺寸
            if scale:
                new_size = (int(video.size[0] * scale), int(video.size[1] * scale))
            elif width and height:
                new_size = (width, height)
            elif width:
                aspect_ratio = video.size[0] / video.size[1]
                new_size = (width, int(width / aspect_ratio))
            elif height:
                aspect_ratio = video.size[0] / video.size[1]
                new_size = (int(height * aspect_ratio), height)
            else:
                return input_path
            
            # 调整尺寸
            resized_video = video.resize(new_size)
            
            # 保存
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            resized_video.write_videofile(output_path)
            
            video.close()
            resized_video.close()
            
            logger.info(f"视频尺寸调整成功: {video.size} -> {new_size}")
            return output_path
            
        except Exception as e:
            logger.error(f"视频尺寸调整失败: {e}")
            raise
    
    def change_speed(self, input_path: str, output_path: str, speed_factor: float) -> str:
        """改变视频速度"""
        try:
            video = mp.VideoFileClip(input_path)
            
            # 改变速度
            speed_video = video.speedx(speed_factor)
            
            # 保存
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            speed_video.write_videofile(output_path)
            
            video.close()
            speed_video.close()
            
            logger.info(f"视频速度调整成功: {speed_factor}x")
            return output_path
            
        except Exception as e:
            logger.error(f"视频速度调整失败: {e}")
            raise
    
    def add_audio(self, video_path: str, audio_path: str, output_path: str) -> str:
        """添加音频到视频"""
        try:
            video = mp.VideoFileClip(video_path)
            audio = mp.AudioFileClip(audio_path)
            
            # 添加音频
            video_with_audio = video.set_audio(audio)
            
            # 保存
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            video_with_audio.write_videofile(output_path)
            
            video.close()
            audio.close()
            video_with_audio.close()
            
            logger.info(f"音频添加成功: {audio_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"音频添加失败: {e}")
            raise
    
    def extract_audio(self, video_path: str, output_path: str) -> str:
        """从视频提取音频"""
        try:
            video = mp.VideoFileClip(video_path)
            
            if video.audio is None:
                raise ValueError("视频没有音频轨道")
            
            # 提取音频
            audio = video.audio
            
            # 保存音频
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            audio.write_audiofile(output_path)
            
            video.close()
            audio.close()
            
            logger.info(f"音频提取成功: {output_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"音频提取失败: {e}")
            raise
    
    def batch_process(self, input_dir: str, output_dir: str, 
                     operations: List[str], **kwargs) -> List[str]:
        """批量处理视频"""
        try:
            # 获取所有视频文件
            video_files = []
            for format in self.supported_formats:
                video_files.extend(Path(input_dir).glob(f"*{format}"))
                video_files.extend(Path(input_dir).glob(f"*{format.upper()}"))
            
            if not video_files:
                logger.warning(f"在 {input_dir} 中没有找到视频文件")
                return []
            
            # 创建输出目录
            Path(output_dir).mkdir(parents=True, exist_ok=True)
            
            results = []
            for i, video_file in enumerate(video_files):
                logger.info(f"处理视频 {i+1}/{len(video_files)}: {video_file.name}")
                
                output_path = Path(output_dir) / f"processed_{video_file.stem}.mp4"
                
                try:
                    # 执行操作
                    if "resize" in operations:
                        self.resize_video(str(video_file), str(output_path), **kwargs)
                    elif "speed" in operations:
                        self.change_speed(str(video_file), str(output_path), **kwargs)
                    elif "convert" in operations:
                        self.convert_format(str(video_file), str(output_path), **kwargs)
                    else:
                        # 默认复制
                        import shutil
                        shutil.copy2(str(video_file), str(output_path))
                    
                    results.append(str(output_path))
                    
                except Exception as e:
                    logger.error(f"处理视频 {video_file.name} 失败: {e}")
                    continue
            
            logger.info(f"批量处理完成，成功处理 {len(results)} 个视频")
            return results
            
        except Exception as e:
            logger.error(f"批量处理失败: {e}")
            raise
    
    def create_thumbnail(self, video_path: str, output_path: str, 
                        time_position: float = 0) -> str:
        """创建视频缩略图"""
        try:
            video = mp.VideoFileClip(video_path)
            
            # 获取指定时间的帧
            frame = video.get_frame(time_position)
            
            # 保存为图片
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            Image.fromarray(frame).save(output_path)
            
            video.close()
            
            logger.info(f"缩略图创建成功: {output_path}")
            return output_path
            
        except Exception as e:
            logger.error(f"缩略图创建失败: {e}")
            raise
    
    def get_video_metadata(self, video_path: str) -> Dict[str, Any]:
        """获取视频元数据"""
        try:
            info = self.get_video_info(video_path)
            
            # 使用moviepy获取更多信息
            video = mp.VideoFileClip(video_path)
            
            metadata = {
                **info,
                'has_audio': video.audio is not None,
                'audio_fps': video.audio.fps if video.audio else None,
                'audio_duration': video.audio.duration if video.audio else None,
                'file_size': Path(video_path).stat().st_size,
                'creation_time': Path(video_path).stat().st_ctime
            }
            
            video.close()
            return metadata
            
        except Exception as e:
            logger.error(f"获取视频元数据失败: {e}")
            raise


class VideoAnalyzer:
    """视频分析器"""
    
    def __init__(self):
        self.processor = VideoProcessor()
    
    def analyze_scene_changes(self, video_path: str, threshold: float = 30.0) -> List[float]:
        """分析场景变化"""
        try:
            frames = self.processor.read_video_frames(video_path)
            
            scene_changes = []
            prev_frame = None
            
            for i, frame in enumerate(frames):
                if prev_frame is not None:
                    # 计算帧间差异
                    diff = cv2.absdiff(frame, prev_frame)
                    mean_diff = np.mean(diff)
                    
                    if mean_diff > threshold:
                        # 计算时间戳
                        timestamp = i / self.processor.get_video_info(video_path)['fps']
                        scene_changes.append(timestamp)
                
                prev_frame = frame
            
            logger.info(f"检测到 {len(scene_changes)} 个场景变化")
            return scene_changes
            
        except Exception as e:
            logger.error(f"场景变化分析失败: {e}")
            raise
    
    def detect_motion(self, video_path: str, min_area: int = 500) -> List[Dict[str, Any]]:
        """检测运动"""
        try:
            frames = self.processor.read_video_frames(video_path)
            
            motion_events = []
            prev_frame = None
            
            for i, frame in enumerate(frames):
                if prev_frame is not None:
                    # 转换为灰度
                    gray_prev = cv2.cvtColor(prev_frame, cv2.COLOR_RGB2GRAY)
                    gray_curr = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
                    
                    # 计算光流
                    flow = cv2.calcOpticalFlowFarneback(
                        gray_prev, gray_curr, None, 0.5, 3, 15, 3, 5, 1.2, 0
                    )
                    
                    # 计算运动强度
                    magnitude = np.sqrt(flow[..., 0]**2 + flow[..., 1]**2)
                    mean_magnitude = np.mean(magnitude)
                    
                    if mean_magnitude > 2.0:  # 运动阈值
                        timestamp = i / self.processor.get_video_info(video_path)['fps']
                        motion_events.append({
                            'timestamp': timestamp,
                            'motion_intensity': mean_magnitude,
                            'frame_index': i
                        })
                
                prev_frame = frame
            
            logger.info(f"检测到 {len(motion_events)} 个运动事件")
            return motion_events
            
        except Exception as e:
            logger.error(f"运动检测失败: {e}")
            raise
    
    def analyze_color_distribution(self, video_path: str, sample_frames: int = 10) -> Dict[str, Any]:
        """分析颜色分布"""
        try:
            video_info = self.processor.get_video_info(video_path)
            total_frames = video_info['frame_count']
            
            # 采样帧
            frame_indices = np.linspace(0, total_frames-1, sample_frames, dtype=int)
            frames = self.processor.read_video_frames(video_path)
            
            color_stats = {
                'red_mean': [], 'red_std': [],
                'green_mean': [], 'green_std': [],
                'blue_mean': [], 'blue_std': []
            }
            
            for idx in frame_indices:
                if idx < len(frames):
                    frame = frames[idx]
                    
                    # 计算RGB通道统计
                    for i, color in enumerate(['red', 'green', 'blue']):
                        channel = frame[:, :, i]
                        color_stats[f'{color}_mean'].append(np.mean(channel))
                        color_stats[f'{color}_std'].append(np.std(channel))
            
            # 计算总体统计
            result = {}
            for key, values in color_stats.items():
                result[f'{key}_overall'] = np.mean(values)
                result[f'{key}_trend'] = np.polyfit(range(len(values)), values, 1)[0]
            
            logger.info("颜色分布分析完成")
            return result
            
        except Exception as e:
            logger.error(f"颜色分布分析失败: {e}")
            raise


if __name__ == "__main__":
    # 测试代码
    processor = VideoProcessor()
    analyzer = VideoAnalyzer()
    
    # 测试视频处理
    test_video = "data/input/test_video.mp4"
    if Path(test_video).exists():
        # 获取视频信息
        info = processor.get_video_info(test_video)
        print(f"视频信息: {info}")
        
        # 创建缩略图
        thumbnail_path = processor.create_thumbnail(test_video, "data/output/thumbnail.jpg")
        print(f"缩略图: {thumbnail_path}")
        
        # 分析场景变化
        scene_changes = analyzer.analyze_scene_changes(test_video)
        print(f"场景变化: {scene_changes}")
        
    else:
        print("测试视频不存在，请先准备测试视频") 